#!/usr/bin/env python

import rospy,time,math
from std_msgs.msg import Int32 ,Float64
from geometry_msgs.msg import PoseStamped,Twist,TwistStamped
from mavros_msgs.msg import AttitudeTarget, PositionTarget, State
from mavros_msgs.srv import CommandBool, SetMode
import tf.transformations as tf
import numpy as np
# from control import care,place

class ROS_INIT():
    def ros_init(self):
        rospy.init_node('uav_formation', anonymous=True)
        self.rate = rospy.Rate(self.h)   # Hz，控制节点主循环的执行频率

        # 发布指令
        self.uav_target_pose_pub = [None]*self.uav_num
        for nm in range(self.uav_num):
            self.uav_target_pose_pub[nm] = rospy.Publisher('/uav'+str(nm+1)+'/mavros/setpoint_raw/local', PositionTarget, queue_size=1)
        
        # 数据保存
        # self.error_pub = rospy.Publisher('/error', PoseStamped, queue_size=1)
        # self.uav_vel_pub = [None]*self.uav_num
        # for nm in range(self.uav_num):
        #     self.uav_vel_pub[nm] = rospy.Publisher('/uav'+str(nm+1)+'vel', Float64, queue_size=1)

        # 订阅各无人机位置角度信息
        self.pose_sub = [None]*self.uav_num
        self.pose_sub_vel = [None]*self.uav_num
        for nm in range(self.uav_num):
            self.pose_sub[nm] = rospy.Subscriber('/uav'+str(nm+1)+'/mavros/local_position/pose', PoseStamped, self.pose_callback,queue_size=1,callback_args=nm+1)
        
        # 订阅各无人机速度信息
        self.vel_sub = [None]*self.uav_num  # 速度订阅器列表
        for nm in range(self.uav_num):
            self.vel_sub[nm] = rospy.Subscriber(
                '/uav'+str(nm+1)+'/mavros/local_position/velocity_local',  # PX4本地速度话题
                TwistStamped,  # 速度消息类型（包含线性和角速度）
                self.vel_callback,  # 回调函数
                queue_size=1,
                callback_args=nm+1  # 传递无人机ID（1-based）
            )
        # 订阅目标位置信息
        self.tar_pos_sub = rospy.Subscriber('/uav1/target_pos', PoseStamped, self.tar_pos_callback,queue_size=1)
        # 订阅px4信息
        rospy.Subscriber('/uav'+str(self.uav_id)+'/mavros/state', State, self.uav_px4state_cb, queue_size=1)
        
        # 上锁与解锁服务
        self.uav_armService = rospy.ServiceProxy('/uav'+str(self.uav_id)+'/mavros/cmd/arming', CommandBool)

        # 订阅地面站指令
        rospy.Subscriber('/command', Int32, self.command_callback)

    def pose_callback(self,msg,nm):
        k = nm-1
        self.uav_pose[k][0] = msg.pose.position.x
        self.uav_pose[k][1] = msg.pose.position.y
        self.uav_pose[k][2] = msg.pose.position.z
        local_angles = tf.euler_from_quaternion([msg.pose.orientation.x,msg.pose.orientation.y,msg.pose.orientation.z,msg.pose.orientation.w]) #获取自己的欧拉角
        self.uav_pose[k][3] = local_angles[0]      # 滚转
        self.uav_pose[k][4] = local_angles[1]      # 俯仰
        self.uav_pose[k][5] = local_angles[2]      # 偏航
    def vel_callback(self,msg,nm):
        k = nm-1
        self.uav_vel[k][0] = msg.twist.linear.x
        self.uav_vel[k][1] = msg.twist.linear.y
        self.uav_vel[k][2] = msg.twist.linear.z
    def tar_pos_callback(self,msg):
        self.tar_pos[0] = msg.pose.position.x
        self.tar_pos[1] = msg.pose.position.y
        self.tar_pos[2] = msg.pose.position.z
    def uav_px4state_cb(self,msg):
        self.px4_state = msg

    def command_callback(self, msg):
        self.command = msg.data
        print('uav',self.uav_id,'Receive command: ', self.command)
        if self.command == 0:
            self.state = 'INIT'
            print('state: ', self.state)
        elif self.command == 1:
            self.state = 'ARM'
            print('state: ', self.state)
        elif self.command == 2:
            self.state = 'TAKEOFF'
            print('state: ', self.state)
        elif self.command == 3:
            self.state = 'LAND'
            print('state: ', self.state)
        elif self.command == 4:
            self.state = 'RESET'
            print('state: ', self.state)
        elif self.command == 5:
            self.state = 'START'
            print('state: ', self.state)
        elif self.command == 6:
            self.state = 'UGV_RESET'
            print('state: ', self.state)
        elif self.command == 7:
            self.state = 'WAITING'
            print('state: ', self.state)
        elif self.command == 8:  # 中制导打击状态
            self.state = 'ATTACK'
            print('state: ', self.state)
        else:
            print('state: ', self.state, ' 保持上一个状态')
            
class Controller(ROS_INIT):
    def __init__(self):
        ###################### 读取系统参数 #############################
        self.uav_id = rospy.get_param('uav_id', 1)    # 获取参数 -- 本机id
        print("本机是 {} 号飞机".format(self.uav_id))
        self.uav_num = rospy.get_param('uav_num', 1)    # 获取参数 -- 集结的无人机数量
        print("集结的无人机数量是 {} ".format(self.uav_num))

        ################## 打击功能参数 ########################
        self.pre_defined_altitude = 100.0  # 预设起飞高度（米）
        self.altitude_threshold = 2.0     # 高度误差阈值（米）
        self.target_uav_id = rospy.get_param(f'/uav{self.uav_id}/target_uav_id', 2) # 获取目标无人机ID
        self.attack_threshold = 100.0  # 中制导打击阈值（米）
        self.attack_miss = 2.0  # 打击半径（米）
        self.max_load = 500.0  # 最大过载限幅
        self.attack_gain = 0.8         # 打击控制增益
        self.attack_history_len = 10   # 轨迹历史长度
        self.target_history = []       # 目标轨迹历史
        self.attack_complete = False   # 打击完成标志

        ################## INIT ############
        self.h = 20  # hz，控制节点主循环的执行频率
        self.dt = 1.0/self.h # 时钟步长
        self.uav_pose = [[0,0,0,0,0,0] for i in range(self.uav_num)]        # n*6矩阵表示每个无人机的x.y,z
        self.tar_pos = [0,0,0]                                              # 目标位置    
        self.uav_vel = [[0,0,0] for i in range(self.uav_num)]               # n*6矩阵表示每个无人机的vx,vy,vz
        self.px4_state = State
        self.auto_arm = True   # 自动解锁
        self.fly_alt = [1.0, 0.8]
        # self.zt = np.zeros(((2+4+2+1+4)*self.uav_num)).reshape(-1,1)     
        self.v_max = 15#定义无人机最大速度
        self.r = 2
        self.i_num = 0
        self.pre_pos = [0,0]
        # print(self.zt)       
        

        ################# ROS init ########################        
        self.command = 0     
        self.state = None   
        self.ros_init()
        # rospy.Timer(rospy.Duration(1/10), self.state_timer_callback, oneshot=False)      # 状态机更新state用

    def main(self):
        start_time = rospy.get_time()
        takeoff_pos = [0,0,0]
        land_pos = [0,0,0]
        id = self.uav_id-1
        
        if False:
            print("ERROR!!!")
        else:
            while not rospy.is_shutdown():
                t = rospy.get_time() - start_time
                if self.state == 'WAITING':         # 初始化
                    start_time = rospy.get_time()
                    # ugv_cmd_vel = Twist()
                    # ugv_cmd_vel.linear.x = 0
                    # ugv_cmd_vel.linear.y = 0
                    # self.uav_target_pose_pub[self.uav_id-1].publish(ugv_cmd_vel)
                
                elif self.state == 'ARM':       # 自动解锁              如果当前处于offboard模式会直接起飞
                    # 解锁，auto_arm设置是否可以自动解锁  
                    if((not self.px4_state.armed) and self.auto_arm):  #仿真中可以自动解锁
                        if self.uav_arm():
                            # print(id,self.uav_pose[id])
                            takeoff_pos = [self.uav_pose[id][0], self.uav_pose[id][1], self.fly_alt[0]]
                            self.state = 'WAITING'
                            # time.sleep(0.5)
                            # self.state = 'TAKEOFF'

                elif self.state == 'TAKEOFF':   # 原地起飞    
                    # print('uav起飞')
                    self.pubPxyPzCmd(takeoff_pos, id)
                    # time.sleep(0.5)
                    if self.dis(takeoff_pos,self.uav_pose[id]) < 0.1 and abs(takeoff_pos[2] - self.uav_pose[id][2]) < 0.1:
                        print('uav',id+1,'已完成起飞!')
                        self.state = 'WAITING'    

                elif self.state == 'RESET':
                    init_pos = self.Formation_init()
                    # init_pos = self.init()
                    reset_pos = [init_pos[id*2],init_pos[id*2+1],self.fly_alt[0]]
                    # print(reset_pos)
                    self.pubPxyPzCmd(reset_pos, id)
                    if self.dis(reset_pos,self.uav_pose[id]) < 0.1 and abs(reset_pos[2] - self.uav_pose[id][2]) < 0.1:
                        print('uav',id+1,'已完成初始化!')
                        self.state = 'WAITING'   
                
                elif self.state == 'START':
                    # if id == 0:
                    #     print(t)
                    self.Control()

                elif self.state == 'LAND':
                    print('uav',id+1,'收到降落指令')
                    land_pos = [self.uav_pose[id][0], self.uav_pose[id][1], 0]
                    # time.sleep(0.5)
                    self.state = 'LAND_2'

                elif self.state == 'LAND_2':
                    if(self.uav_pose[id][2] > 0.25):
                        uav_VxyVz_sp = [0.0, 0.0, -0.2]
                        self.pubPxyVzCmd(land_pos, uav_VxyVz_sp, id)
                    else:
                        uav1_VxyVz_sp = [0.0, 0.0, -0.8]
                        self.pubPxyVzCmd(land_pos, uav1_VxyVz_sp, id)
                    if self.dis(land_pos,self.uav_pose[id]) < 0.1 and abs(land_pos[2] - self.uav_pose[id][2]) < 0.1:
                        print('uav',id+1,'已完成降落!')
                        self.state = 'DISARM'    
                
                elif self.state == 'DISARM':
                    if self.uav_disarm():
                        break
                        
                elif self.state == 'ATTACK':  # 中制导打击状态
                    if not self.attack_complete:
                        self.attack_control()
                    else:
                        # 中制导打击完成，切换到末制导
                        self.guindance_control()            
               
                self.rate.sleep()

    def uav_arm(self):
        if self.uav_armService(True):
            print('uav'+str(self.uav_id)+'--已解锁！')
            return True
        else:
            print('uav'+str(self.uav_id)+'--解锁失败！')
            return False
    def uav_disarm(self):
        if self.uav_armService(False):
            print('uav'+str(self.uav_id)+'--已上锁！')
            return True
        else:
            print('uav'+str(self.uav_id)+'--上锁失败！')
            return False
        
    def pubPxyPzCmd(self, cmd_p, id):
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = 1    # pub in ENU local frame
        msg.type_mask = 0b110111111000  # pub pz+py+px. ignore yaw and yaw rate and ax ay az
        msg.position.x = cmd_p[0]   # pub px
        msg.position.y = cmd_p[1]   # pub py
        msg.position.z = cmd_p[2]   # pub local z altitude setpoint
        self.uav_target_pose_pub[id].publish(msg)

    def pubPxyVzCmd(self, cmd_p, cmd_v, id):
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = 1    # pub in ENU local frame
        msg.type_mask = 0b110111011100  # pub vz+py+px. ignore yaw and yaw rate
        msg.position.x = cmd_p[0]   # pub px
        msg.position.y = cmd_p[1]   # pub py
        msg.velocity.z = cmd_v[2]   # pub vz
        self.uav_target_pose_pub[id].publish(msg)

    def pubVxyPzCmd(self, cmd_v, pz, id):
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = 1    # pub in ENU local frame
        msg.type_mask = 0b110111000011  # pub vxvypz
        msg.velocity.x = cmd_v[0]   # pub vx
        msg.velocity.y = cmd_v[1]   # pub vy
        msg.velocity.z = 0          # pub vz
        msg.position.z = pz         # pub pz
        self.uav_target_pose_pub[id].publish(msg)

    def pubVxyzCmd(self, cmd_v, id):
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = 1    # pub in ENU local frame
        msg.type_mask = 0b110111000111  # pub vxvypz
        msg.velocity.x = cmd_v[0]   # pub vx
        msg.velocity.y = cmd_v[1]   # pub vy
        msg.velocity.z = 0          # pub vz
        self.uav_target_pose_pub[id].publish(msg)

    def Control(self):
        id = self.uav_id-1
        # cmd_vel = [1,1,0]
        k = 1.4
        cmd_vel = [0,0,0]
        vxt = (self.tar_pos[0]-self.pre_pos[0])*self.h
        vyt = (self.tar_pos[1]-self.pre_pos[1])*self.h
        print(vxt,vyt)
        cmd_vel[0] = vxt + k*(self.tar_pos[0] - self.uav_pose[id][0])
        cmd_vel[1] = vyt + k*(self.tar_pos[1] - self.uav_pose[id][1])
        print(self.dis(self.tar_pos,self.uav_pose[id]))
        self.pubVxyzCmd(cmd_vel, self.uav_id-1)

        self.pre_pos = [self.tar_pos[0],self.tar_pos[1]]

    def dis(self,a,b):
        return math.sqrt((a[0]-b[0])**2 + (a[1]-b[1])**2)

    def Formation_init(self):
        # 初始化编队位置（示例：圆形编队）
        radius = 10.0  # 编队半径
        angles = [2 * math.pi * i / self.uav_num for i in range(self.uav_num)]
        positions = []
        for angle in angles:
            x = radius * math.cos(angle)
            y = radius * math.sin(angle)
            positions.extend([x, y])
        return positions

    def attack_control(self):
        """中制导打击控制函数"""
        id = self.uav_id - 1  # 本机索引
        target_id = self.target_uav_id - 1  # 目标机索引

        # 检查目标ID有效性
        if target_id < 0 or target_id >= self.uav_num:
            rospy.logerr("无效的目标无人机ID: %d", self.target_uav_id)
            self.state = 'WAITING'  # 切换到等待状态
            return

        # 获取目标当前位置
        target_pos = self.uav_pose[target_id]
        self_pos = self.uav_pose[id]
        
        # 更新目标轨迹历史
        self.target_history.append([target_pos[0], target_pos[1], target_pos[2]])
        if len(self.target_history) > self.attack_history_len:
            self.target_history.pop(0)
        
        # 计算与目标的距离
        distance = self.dis(target_pos, self_pos)
        
        # 当距离小于阈值时，认为打击完成
        if distance < self.attack_threshold:
            rospy.loginfo("我方ID: %d,已接近目标至 %0.2f 米，中制导完成！", self.uav_id,distance)
            self.attack_complete = True
            return
            
        # 计算目标的平均速度（基于历史轨迹）后续改为数据驱动或机器学习方法
        if len(self.target_history) >= 2:
            dt = 1.0 / self.h * (len(self.target_history) - 1)#相邻采样时间差 Δt0​=1/self.h，总时间差 Δt=n×Δt0​
            avg_vx = (self.target_history[-1][0] - self.target_history[0][0]) / dt
            avg_vy = (self.target_history[-1][1] - self.target_history[0][1]) / dt
            avg_vz = (self.target_history[-1][2] - self.target_history[0][2]) / dt
        else:
            avg_vx, avg_vy, avg_vz = 0.0, 0.0, 0.0
        
        # 预测目标未来位置（提self.attack_miss前量计算）
        lead_time = 2.0  # 预测未来2秒
        predicted_x = target_pos[0] + avg_vx * lead_time
        predicted_y = target_pos[1] + avg_vy * lead_time
        predicted_z = target_pos[2] + avg_vz * lead_time
        
        # 计算到达预测位置所需的速度
        desired_vx = (predicted_x - self_pos[0]) * self.attack_gain
        desired_vy = (predicted_y - self_pos[1]) * self.attack_gain
        desired_vz = (predicted_z - self_pos[2]) * self.attack_gain
        
        # 速度限制
        desired_vx = max(-self.v_max, min(desired_vx, self.v_max))
        desired_vy = max(-self.v_max, min(desired_vy, self.v_max))
        desired_vz = max(-self.v_max, min(desired_vz, self.v_max))
        

        # 起飞阶段：如果当前高度低于预设高度，先上升到预设高度
        if self_pos[2] < self.pre_defined_altitude - self.altitude_threshold:
            # 垂直上升，水平速度为0
            self.pubVxyPzCmd([0.0, 0.0], self.pre_defined_altitude, id)
            rospy.loginfo("起飞阶段：上升至预设高度 %0.2f 米，当前高度 %0.2f 米", 
                            self.pre_defined_altitude, self_pos[2])
            return

        # 高度控制策略：根据距离选择不同的高度控制方式
        if distance > 50:  # 远距离时，保持预设高度
            self.pubVxyPzCmd([desired_vx, desired_vy], self.pre_defined_altitude, id)
            rospy.loginfo("我方ID: %d, 目标ID: %d, 远距离追踪：保持高度 %0.2f 米，距离目标 %0.2f 米",
              self.uav_id, self.target_uav_id, self.pre_defined_altitude, distance)
        # else:  # 近距离时，跟随目标高度
        #     self.pubVxyPzCmd([desired_vx, desired_vy], predicted_z, id)
        #     rospy.loginfo("近距离追踪：跟随目标高度，距离目标 %0.2f 米", distance)

    def guindance_control(self):
        """末制导打击控制函数"""
        id = self.uav_id - 1  # 本机索引
        target_id = self.target_uav_id - 1  # 目标机索引

        # 检查目标ID有效性
        if target_id < 0 or target_id >= self.uav_num:
            rospy.logerr("无效的目标无人机ID: %d", self.target_uav_id)
            self.state = 'WAITING'  # 切换到等待状态
            return
        # 获取当前位置
        target_pos = np.array(self.uav_pose[target_id][:3])  # 取前3个元素：x,y,z
        self_pos = np.array(self.uav_pose[id][:3])           # 取前3个元素：x,y,z

        # 获取当前速度（转换为numpy数组）
        target_vel = np.array(self.uav_vel[target_id])  # 速度列表转数组
        self_vel = np.array(self.uav_vel[id])           # 速度列表转数组

        # 计算速度模长（避免除以0，加一个极小值）
        self_vel_norm = np.linalg.norm(self_vel) + 1e-6
        target_vel_norm = np.linalg.norm(target_vel) + 1e-6

        # 计算弹目距离（避免除以0）
        R_MT = self.dis(target_pos, self_pos) + 1e-6

        # 计算单位方向向量
        eR = (target_pos - self_pos) / R_MT  
        eVm = self_vel / self_vel_norm
        eVt = target_vel / target_vel_norm

        # 计算Omega_MT
        Omega_MT = np.cross(self_vel - target_vel, eR) / R_MT

        # 三维比例导引(导航比为5)
        a_PNG = 5.0*np.cross(Omega_MT, self_vel)
        # 计算下一时刻速度
        v_next = self_vel + a_PNG*self.dt
        # 速度限幅
        v_next = v_next/np.linalg.norm(v_next)*self_vel_norm
        # 
        desired_vx = v_next[0]
        desired_vy = v_next[1]
        desired_vz = v_next[2]

        # 二维比例导引
        # # 计算弹道倾角
        # target_gamma = np.arctan2(target_vel[1], target_vel[0])
        # self_gamma = np.arctan2(self_vel[1], self_vel[0])

        # # 计算弹目视线角
        # q_MT = np.arctan2(target_pos[1]-self_pos[1], target_pos[0]-self_pos[0])
        
        # # 计算与目标的距离
        # R_MT = self.dis(target_pos, self_pos)
        # dq = (target_vel_norm*np.sin(target_gamma-q_MT)-self_vel_norm*np.sin(self_gamma-q_MT))/R_MT

        # # 计算比例导引
        # a = 3.0*self_vel_norm*dq
        # # 过载限幅
        # if a > self.max_load:
        #     a = self.max_load
        # elif a < -self.max_load:
        #     a = -self.max_load
        # self_gamma_next = self_gamma + a*self.dt
        # # 计算所需的速度
        # desired_vx = self_vel_norm*np.cos(self_gamma_next)
        # desired_vy = self_vel_norm*np.sin(self_gamma_next)
        # desired_vz = 0.0
        
        # # 速度限制
        # desired_vx = max(-self.v_max, min(desired_vx, self.v_max))
        # desired_vy = max(-self.v_max, min(desired_vy, self.v_max))
        # desired_vz = max(-self.v_max, min(desired_vz, self.v_max))
        
        # 当距离小于阈值时，认为打击完成
        if R_MT < self.attack_miss:
            self.pubVxyPzCmd([0, 0], self.pre_defined_altitude, id)            
            rospy.loginfo("已接近目标至 %0.2f 米，末制导完成！", R_MT)
        else:
            self.pubVxyPzCmd([desired_vx, desired_vy], self.pre_defined_altitude, id)
            rospy.loginfo("我方ID: %d, 目标ID: %d, 末制导打击：保持高度 %0.2f 米，距离目标 %0.2f 米",
              self.uav_id, self.target_uav_id, self.pre_defined_altitude, R_MT)
      

if __name__ == "__main__":
    formation = Controller()
    try:
        formation.main()
    except rospy.ROSInterruptException:
        pass
